import sys
import os

import requests


class BaseNode:
    def __init__(self):
        self.type = None
        self.required: dict = {}

    # 返回类型
    RETURN_TYPES = None

    # 节点执行方法
    FUNCTION = "execute"

    # 所属类别
    CATEGORY = None

    @classmethod
    def INPUT_TYPES(cls):
        return {
        }

    def execute(self):
        return None


class UserInput(BaseNode):
    """
    提示词输入
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "text": ("STRING", {
                    "multiline": True,
                    "default": "Your Input"
                }),
            }
        }

    RETURN_TYPES = ("STRING",)
    FUNCTION = "execute"
    CATEGORY = "AutoGen"

    def execute(self, text):
        return (text,)


class SessionDetection:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "INPUT": ("STRING", {"forceInput": True})
            }
        }

    RETURN_TYPES = ("STRING",)
    FUNCTION = "execute"
    CATEGORY = "AutoGen"

    def translate_prompt(self, INPUT):
        url = "http://llm7.shauetech.com:8004/v1/chat/completions"
        data = {
            "stream": False,
            "model": "Qwen2-72b-Chat",
            "messages": [
                {
                    "role": "system",
                    "content": "我现在有一段对话，你是一个情感检测专家，检测这段对话，如果你检测到有冲突发生的可能性你就发送：“哔哔哔”，否则发送：相安无事"
                },
                {
                    "role": "user",
                    "content": "这段对话为: " + INPUT
                }

            ],
            "temperature": 0,
            "top_p": 1,
            "n": 1,
            "max_tokens": 2000,
            "presence_penalty": 0,
            "frequency_penalty": 0
        }
        response_ = requests.post(url, json=data, timeout=3000)
        answer_response = response_.json()
        print(answer_response)
        try:
            answer = answer_response["choices"][0]['message']['content']
        except Exception as e:
            answer = "检测失败了！！！"
        return answer

    def execute(self, INPUT):
        prompt_polish_text = self.translate_prompt(INPUT)
        return (prompt_polish_text,)


NODE_CLASS_MAPPINGS = {
    "SessionDetection": SessionDetection,
    "UserInput": UserInput
}
NODE_DISPLAY_NAME_MAPPINGS = {
    "SessionDetection": "SessionDetection"
}
